import load2
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import q1
import keras
import numpy as np
import pickle

allTaskJing=q1.allJingA+q1.allJingB
allTaskWei=q1.allWeiA+q1.allWeiB

allPos=[]
allXiane=[]
for i in range(1,1878):
    if load2.getJingWei(i):
        allPos.append(load2.getJingWei(i))
        allXiane.append(load2.getCell(i,load2.xiane))

estimator = KMeans(n_clusters=3) # 构造聚类器
estimator.fit(allPos) # 聚类

C1Member=[]
C2Member=[]
C3Member=[]
for i in range(len(allXiane)):
    if estimator.labels_[i]==0:
        C1Member.append(allXiane[i])
    elif estimator.labels_[i]==1:
        C2Member.append(allXiane[i])
    else:
        C3Member.append(allXiane[i])

print(estimator.cluster_centers_)

def dist(jing1,wei1,jing2,wei2):
    return (jing1-jing2)**2+(wei1-wei2)**2

# 按聚集点给任务分类
def classTask(allTaskJing,allTaskWei):
    C1TaskJing = []
    C2TaskJing = []
    C3TaskJing = []
    C1TaskWei = []
    C2TaskWei = []
    C3TaskWei = []
    allDist = []
    for i in range(len(allTaskJing)):
        d1 = dist(allTaskJing[i], allTaskWei[i], estimator.cluster_centers_[0, 0], estimator.cluster_centers_[0, 1])
        d2 = dist(allTaskJing[i], allTaskWei[i], estimator.cluster_centers_[1, 0], estimator.cluster_centers_[1, 1])
        d3 = dist(allTaskJing[i], allTaskWei[i], estimator.cluster_centers_[2, 0], estimator.cluster_centers_[2, 1])
        d, sub = min([(d1, 0), (d2, 1), (d3, 2)], key=lambda x: x[0])
        allDist.append(d)
        if sub == 0:
            C1TaskJing.append(allTaskJing[i])
            C1TaskWei.append(allTaskWei[i])
        elif sub == 1:
            C2TaskJing.append(allTaskJing[i])
            C2TaskWei.append(allTaskWei[i])
        else:
            C3TaskJing.append(allTaskJing[i])
            C3TaskWei.append(allTaskWei[i])
    return C1TaskJing,C2TaskJing,C3TaskJing,C1TaskWei,C2TaskWei,C3TaskWei,allDist

C1TaskJing,C2TaskJing,C3TaskJing,C1TaskWei,C2TaskWei,C3TaskWei,allDist=classTask(q1.allJingA,q1.allWeiA)

if __name__=='__main__':
    # 训练定价模型
    model = keras.Sequential()
    model.add(keras.layers.Dense(1, input_shape=(1,), activation='linear'))
    model.compile(optimizer='adam', loss='mse')
    allDist=np.array(allDist)
    print(allDist)
    allJiageA=np.array(q1.allJiageA)
    allJiageA=np.sqrt(allJiageA)
    model.fit(allDist, allJiageA, epochs=5000, batch_size=250)
    pickle.dump(model, open('q3dingjia.pkl', 'wb'), protocol=2)

    print(model.predict(allDist)**2)

    jing=[i[0] for i in allPos]
    wei=[i[1] for i in allPos]

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(jing,wei,c='black')
    ax.scatter(C1TaskJing,C1TaskWei,c='red')
    ax.scatter(C2TaskJing,C2TaskWei,c='green')
    ax.scatter(C3TaskJing,C3TaskWei,c='blue')
    for i in estimator.cluster_centers_:
        ax.scatter(i[0], i[1], c='gray')
    plt.show()
